import torch.nn as nn
import torch.nn.init as init

from .nets.backbone_空间注意力 import DCLC
from .nets.junction_decoder import SuperpointDecoder
from .nets.heatmap_decoder import HeatmapDecoder
from .nets.descriptor_decoder import SuperpointDescriptor


def get_model(model_cfg=None, loss_weights=None, loss_funcs=None, mode="train"):
    """ Get model based on the model configuration. """
    # Check dataset config is given
    if model_cfg is None:
        raise ValueError("[Error] The model config is required!")

    # List the supported options here
    print("\n\n\t--------Initializing model----------")
    supported_arch = ["simple"]
    if not model_cfg["model_architecture"] in supported_arch:
        raise ValueError(
            "[Error] The model architecture is not in supported arch!")

    if model_cfg["model_architecture"] == "simple":
        #SOLD2Net网络模型类的实例化对象
        #出入配置python字典，初始化网络参数
        model = DCLCL(model_cfg, loss_weights, loss_funcs)
    else:
        raise ValueError(
            "[Error] The model architecture is not in supported arch!")

    # Optionally register loss weights to the model
    if mode == "train":
        if loss_weights is not None:
            for param_name, param in loss_weights.items():
                #判断对象是否是一个已知的类型
                if isinstance(param, nn.Parameter):
                    print("\t [Debug] Adding %s with value %f to model"
                          % (param_name, param.item()))
                    #向网络添加parameter
                    model.register_parameter(param_name, param)
        else:
            raise ValueError(
                "[Error] the loss weights can not be None in dynamic weighting mode during training.")

    # Display some summary info.
    print("\tModel architecture: %s" % model_cfg["model_architecture"])     #simple
    print("\tBackbone: %s" % model_cfg["backbone"])
    print("\tJunction decoder: %s" % model_cfg["junction_decoder"])         #superpoint_decoder
    print("\tHeatmap decoder: %s" % model_cfg["heatmap_decoder"])
    print("\t-------------------------------------")

    #网络模型的类的实例化参数
    return model

#输入图像，返回一个python字典包括"junctions" "heatmap" "descriptors"，分别为处理得到的结果,tensor变量
#过程:
# 图像传入，backbone网络的实例化对象进行处理输出特征图变量
#特征图分别传入SuperPointDecoder、heatmapDecoder、DescriptorDecoder的类实例化对象，分别得到相应的结果,tensor变量
class DCLCL(nn.Module):
    """ Full network for SOLD². """
    def __init__(self, model_cfg, loss_weights, loss_funcs):
        super(DCLCL, self).__init__()
        self.name = model_cfg["model_name"]             # DCLCL
        self.cfg = model_cfg

        # List supported network options
        self.supported_backbone = ["DCLC"]
        # backbone网络模型的类的实例化对象
        self.backbone_net = self.get_backbone()

        # List supported junction decoder options
        self.supported_junction_decoder = ["superpoint_decoder"]
        #Hourglass backbone输入到SuperPoint解码器的类对象 输出通道为65 大小为原图的1/8
        self.junction_decoder = self.get_junction_decoder()

        # List supported heatmap decoder options
        self.supported_heatmap_decoder = ["heatmap_decoder"]
        #Hourglass backbone输入，pixel_shuffle解码方式PixelShuffleDecoder解码器的类对象 输出通道为2 大小与原图一致
        self.heatmap_decoder = self.get_heatmap_decoder()

        # List supported descriptor decoder options
        if "descriptor_decoder" in self.cfg:
            self.supported_descriptor_decoder = ["superpoint_descriptor"]
        #Hourglass backbone输入，SuperpointDescriptor解码器的类对象 输出通道为128 大小为原图1/4
            self.descriptor_decoder = self.get_descriptor_decoder()

        self.loss_func = None
        if loss_funcs is not None:
            # Define the total loss
            from .loss import TotalLoss
            policy = model_cfg.get("weighting_policy", "static")
            self.loss_func = TotalLoss(loss_funcs, loss_weights, policy).cuda()

        # Initialize the model weights
        #apply(fn)将函数fn递归地应用到模块自身以及该模块的每一个子模块，该方法通常用来初始化一个模型中的参数
        self.apply(weight_init)

    # 继承的父类nn.Moudle使用类+参数的语法直接调用__call__(self)函数，__call__(self)函数中会返回forward函数的结果
    # 因此，forward函数可以直接通过类名被调用，而不用实例化对象
    def forward(self, input_images, junc_target=None, heatmap_target=None, angle_target=None, valid_mask=None):
        # The backbone
        # features_coarse -- H/4 x W/4 x 256    1/4
        # features_fine -- H x W x 64     1
        features_fine, features_coarse, angle0, angle1, angle2 = self.backbone_net(input_images)


        # junction decoder
        junctions = self.junction_decoder(features_coarse)

        # heatmap decoder
        heatmaps = self.heatmap_decoder(features_fine)

        outputs = {"junctions": junctions, "heatmap": heatmaps,
                   "angle0": angle0, "angle1": angle1, "angle2": angle2}

        # Descriptor decoder
        if "descriptor_decoder" in self.cfg:
            outputs["descriptors"] = self.descriptor_decoder(features_coarse)

        # 如果有真值，则计算损失
        if self.loss_func is not None and junc_target is not None and heatmap_target is not None:
            loss_output = self.loss_func(junctions, junc_target, heatmaps, heatmap_target, angle0,
                                         angle1, angle2, angle_target, valid_mask)
            outputs["loss"] = loss_output

        return outputs

    def forward_descriptors(self, input_images, input_images2, junc_target=None, junc_target2=None, heatmap_target=None, heatmap_target2=None,
                            line_points=None, line_points2=None, line_indices=None, angle_target=None, angle_target2=None, epoch=None, valid_mask=None, valid_mask2=None):
        # Run the forward pass
        outputs = self.forward(input_images)
        outputs2 = self.forward(input_images2)

        # Compute losses
        losses = self.loss_func.forward_descriptors(
            outputs["junctions"], outputs2["junctions"],
            junc_target, junc_target2, outputs["heatmap"], outputs2["heatmap"],
            heatmap_target, heatmap_target2, line_points, line_points2,
            line_indices, outputs['descriptors'], outputs2['descriptors'],
            outputs["angle0"],  outputs["angle1"],  outputs["angle2"],
            outputs2["angle0"],  outputs2["angle1"],  outputs2["angle2"],
            angle_target, angle_target2,
            epoch, valid_mask, valid_mask2)


        return outputs, outputs2, losses

    def get_backbone(self):
        """ Retrieve the backbone encoder network. """
        if not self.cfg["backbone"] in self.supported_backbone:
            raise ValueError(
                "[Error] The backbone selection is not supported.")

        # lcnn backbone (stacked hourglass)
        if self.cfg["backbone"] == "DCLC":
            backbone_cfg = self.cfg["backbone_cfg"]
            #**backbone_cfg 将字典解开成独立的元素作为形参
            backbone = DCLC(**backbone_cfg)
        else:
            raise ValueError(
                "[Error] The backbone selection is not supported.")
        #backbone网络模型的类的实例化对象，feat_channel通道数
        return backbone

    def get_junction_decoder(self):
        """ Get the junction decoder. """
        if (not self.cfg["junction_decoder"]
            in self.supported_junction_decoder):
            raise ValueError(
                "[Error] The junction decoder selection is not supported.")

        # superpoint decoder
        if self.cfg["junction_decoder"] == "superpoint_decoder":
            decoder = SuperpointDecoder(256,
                                        self.cfg["backbone"])
        else:
            raise ValueError(
                "[Error] The junction decoder selection is not supported.")

        return decoder

    def get_heatmap_decoder(self):
        """ Get the heatmap decoder. """
        if not self.cfg["heatmap_decoder"] in self.supported_heatmap_decoder:
            raise ValueError(
                "[Error] The heatmap decoder selection is not supported.")

        # Pixel_shuffle decoder
        if self.cfg["heatmap_decoder"] == "heatmap_decoder":
            if self.cfg["backbone"] == "DCLC":
                decoder = HeatmapDecoder()
            else:
                raise ValueError("[Error] Unknown backbone option.")
        else:
            raise ValueError(
                "[Error] The heatmap decoder selection is not supported.")

        return decoder

    def get_descriptor_decoder(self):
        """ Get the descriptor decoder. """
        if (not self.cfg["descriptor_decoder"]
            in self.supported_descriptor_decoder):
            raise ValueError(
                "[Error] The descriptor decoder selection is not supported.")

        # SuperPoint descriptor
        if self.cfg["descriptor_decoder"] == "superpoint_descriptor":
            decoder = SuperpointDescriptor(256)
        else:
            raise ValueError(
                "[Error] The descriptor decoder selection is not supported.")

        return decoder


def weight_init(m):
    """ Weight initialization function. """
    # Conv2D
    #判断对象是否是一个已知的类型
    if isinstance(m, nn.Conv2d):
        #正态分布初始化
        #xavier保证网络输入与输出的方差相同，即输入的方差不要过小(sigmoid输出为线性)，输入的方差过大(sigmoid输出的梯度消失)
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    # Batchnorm
    elif isinstance(m, nn.BatchNorm2d):
        #初始化参数值符合正态分布
        init.normal_(m.weight.data, mean=1, std=0.02)
        #初始化为常量
        init.constant_(m.bias.data, 0)
    # Linear
    elif isinstance(m, nn.Linear):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    else:
        pass

